// ========================================================================
// ------------------------------------------------------------------------
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at 
// http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ========================================================================

package org.tools.wstt.servlet;

import java.io.IOException;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.ProtocolException;
import java.net.URL;
import java.net.URLConnection;
import java.util.Enumeration;
import java.util.GregorianCalendar;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;

import javax.servlet.Servlet;
import javax.servlet.ServletConfig;
import javax.servlet.ServletContext;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.apache.commons.io.IOUtils;
import org.apache.commons.lang.StringUtils;
import org.tools.wstt.reference.ApplicationComponent;
import org.tools.wstt.utils.logger.IWsttLogger;
import org.tools.wstt.utils.logger.WsttLoggerFactory;

/**
 * Proxy servlet.
 * 
 * @author jlamande : this version behaves not only like a real proxy (CONNECT)
 *         but like a request-response forwarder with abilities of customization
 *         : configuration of destination
 */
public class ProxyServlet implements Servlet, ApplicationComponent {

	protected String destProtocol = "http";
	protected String destHost;
	protected int destPort;
	protected String httpVersion;
	protected IWsttLogger mLogger;

	protected HashSet _DontProxyHeaders = new HashSet();
	{
		_DontProxyHeaders.add("proxy-connection");
		_DontProxyHeaders.add("connection");
		_DontProxyHeaders.add("keep-alive");
		_DontProxyHeaders.add("transfer-encoding");
		_DontProxyHeaders.add("te");
		_DontProxyHeaders.add("trailer");
		_DontProxyHeaders.add("proxy-authorization");
		_DontProxyHeaders.add("proxy-authenticate");
		_DontProxyHeaders.add("upgrade");
	}

	private ServletConfig config;
	private ServletContext mContext;

	/*
	 * (non-Javadoc)
	 * 
	 * @see javax.servlet.Servlet#init(javax.servlet.ServletConfig)
	 */
	public void init(ServletConfig config) throws ServletException {
		this.config = config;
		this.mContext = config.getServletContext();
		this.mLogger = WsttLoggerFactory.getLogger(this.getClass());
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javax.servlet.Servlet#getServletConfig()
	 */
	public ServletConfig getServletConfig() {
		return config;
	}

	private boolean initConfig() {
		destHost = (String) mContext.getAttribute(APP_CONFIG_PROXY_DEST_HOST);
		String vDestPort = (String) mContext
				.getAttribute(APP_CONFIG_PROXY_DEST_PORT);
		destPort = (!StringUtils.isEmpty(vDestPort)) ? Integer
				.parseInt(vDestPort) : 80;
		return !StringUtils.isEmpty(destHost);
	}

	/**
	 * 
	 * @see javax.servlet.Servlet#service(javax.servlet.ServletRequest,
	 *      javax.servlet.ServletResponse)
	 */
	public void service(ServletRequest pRequest, ServletResponse pResponse)
			throws ServletException, IOException {
		long vStartTime = System.currentTimeMillis();
		// use proxy configuration
		initConfig();
		HttpURLConnection vConnection = null;
		try {
			HttpServletRequest vRequest = (HttpServletRequest) pRequest;
			HttpServletResponse vResponse = (HttpServletResponse) pResponse;

			// create the connection
			vConnection = getURLConnection(vRequest);
			// prepare request
			boolean vHasContent = processRequestHeaders(vRequest, vConnection);

			vConnection.setDoInput(true);

			// use request content if any
			if (vHasContent) {
				InputStream vIn = vRequest.getInputStream();
				vConnection.setDoOutput(true);
				IOUtils.copy(vIn, vConnection.getOutputStream());
			}

			// Connect
			vConnection.connect();

			// handler status codes etc.
			mLogger.info("response = " + vConnection.getResponseCode());
			vResponse.setStatus(vConnection.getResponseCode());

			// Process Response Headers
			processResponseHeaders(vConnection, vResponse);

			// 
			InputStream vProxyIn = vConnection.getErrorStream();
			if (vProxyIn == null) {
				try {
					vProxyIn = vConnection.getInputStream();
				} catch (Exception vE) {
					mLogger.info("stream", vE);
					vProxyIn = vConnection.getErrorStream();
				}
			}

			// Handle
			if (vProxyIn != null)
				IOUtils.copy(vProxyIn, vResponse.getOutputStream());

			traceCall(vConnection.getURL().toString(), System
					.currentTimeMillis()
					- vStartTime, "");
		} finally {
			if (vConnection != null) {
				vConnection.disconnect();
			}
		}
	}

	/**
	 * 
	 * @param pRequest
	 * @return
	 */
	protected HttpURLConnection getURLConnection(HttpServletRequest pRequest) {
		HttpURLConnection vConnection = null;

		try {
			URL vUrl = new URL("http", destHost, destPort, getDestUri(pRequest));
			mLogger.info("URL=" + vUrl);
			vConnection = (HttpURLConnection) vUrl.openConnection();
			vConnection.setAllowUserInteraction(false);

			// Set method
			vConnection = (HttpURLConnection) vConnection;
			vConnection.setRequestMethod(pRequest.getMethod());
			vConnection.setInstanceFollowRedirects(false);
		} catch (ProtocolException vE) {
			mLogger.error("Fatal transport error: " + vE.getMessage());
			if (mLogger.isDebugEnabled()) {
				mLogger.debug("Details : ", vE);
			}
		} catch (IOException vE) {
			mLogger.error("Fatal transport error: " + vE.getMessage());
			if (mLogger.isDebugEnabled()) {
				mLogger.debug("Details : ", vE);
			}
		}

		return vConnection;
	}

	/**
	 * 
	 * @param pRequest
	 * @return
	 */
	protected String getDestUri(HttpServletRequest pRequest) {
		String vUri = pRequest.getRequestURI().replaceAll(INTERNAL_PATH_FOR_PROXYING, "");
		if (pRequest.getQueryString() != null)
			vUri += "?" + pRequest.getQueryString();
		return vUri;
	}

	/**
	 * 
	 * @param pRequest
	 * @param pConnection
	 */
	protected boolean processRequestHeaders(HttpServletRequest pRequest,
			URLConnection pConnection) {
		// check connection header
		String vConnectionHdr = pRequest.getHeader("Connection");
		if (vConnectionHdr != null) {
			vConnectionHdr = vConnectionHdr.toLowerCase();
			if (vConnectionHdr.equals("keep-alive")
					|| vConnectionHdr.equals("close"))
				vConnectionHdr = null;
		}

		// copy headers
		boolean vXForwardedFor = false;
		boolean vHasContent = false;
		Enumeration vEnm = pRequest.getHeaderNames();
		while (vEnm.hasMoreElements()) {
			// TODO could be better than this!
			String vHdr = (String) vEnm.nextElement();
			String vLhdr = vHdr.toLowerCase();

			if (_DontProxyHeaders.contains(vLhdr))
				continue;
			if (vConnectionHdr != null && vConnectionHdr.indexOf(vLhdr) >= 0)
				continue;

			if ("content-type".equals(vLhdr))
				vHasContent = true;

			Enumeration vVals = pRequest.getHeaders(vHdr);
			while (vVals.hasMoreElements()) {
				String vVal = (String) vVals.nextElement();
				if (vVal != null) {
					pConnection.addRequestProperty(vHdr, vVal);
					mLogger.info("req " + vHdr + ": " + vVal);
					vXForwardedFor |= "X-Forwarded-For".equalsIgnoreCase(vHdr);
				}
			}
		}

		// Proxy headers
		pConnection.setRequestProperty("Via", "1.1 (jetty)");
		if (!vXForwardedFor)
			pConnection.addRequestProperty("X-Forwarded-For", pRequest
					.getRemoteAddr());

		// a little bit of cache control
		String vCacheControl = pRequest.getHeader("Cache-Control");
		if (vCacheControl != null
				&& (vCacheControl.indexOf("no-cache") >= 0 || vCacheControl
						.indexOf("no-store") >= 0))
			pConnection.setUseCaches(false);

		return vHasContent;
	}

	/**
	 * 
	 * @param pConnection
	 * @param pResponse
	 */
	protected void processResponseHeaders(HttpURLConnection pConnection,
			HttpServletResponse pResponse) {
		// clear response defaults.
		pResponse.setHeader("Date", null);
		pResponse.setHeader("Server", null);

		// set response headers
		int vH = 0;
		String vHdr = pConnection.getHeaderFieldKey(vH);
		String vVal = pConnection.getHeaderField(vH);
		while (vHdr != null || vVal != null) {
			String lhdr = vHdr != null ? vHdr.toLowerCase() : null;
			if (vHdr != null && vVal != null
					&& !_DontProxyHeaders.contains(lhdr))
				pResponse.addHeader(vHdr, vVal);

			mLogger.info("res " + vHdr + ": " + vVal);

			vH++;
			vHdr = pConnection.getHeaderFieldKey(vH);
			vVal = pConnection.getHeaderField(vH);
		}
		pResponse.addHeader("Via", "1.1 (WSTT)");
	}

	/**
	 * 
	 * @param uri
	 * @param pTime
	 */
	protected void traceCall(String uri, long pTime, String pRespBody) {
		Map vMap = new HashMap();
		vMap.put("uri", uri);
		vMap.put("date", GregorianCalendar.getInstance().getTime());
		vMap.put("time", new Long(pTime));
		vMap.put("destination", "http://" + destHost + ":" + destPort);
		vMap.put("response", pRespBody);
		List vList = (List) mContext.getAttribute(APP_PROXY_CALLS);
		if (vList.size() >= APP_PROXY_REMIND_CALLS) {
			vList.remove(0);
		}
		computeStats(pTime);
		vList.add(vMap);
	}

	/**
	 * 
	 * @param pTime
	 */
	protected void computeStats(long pTime) {
		Integer vCount = (Integer) mContext
				.getAttribute(APP_PROXY_CALLS_TOTAL_COUNT);
		vCount = new Integer(vCount.intValue() + 1);
		mContext.setAttribute(APP_PROXY_CALLS_TOTAL_COUNT, vCount);

		Long vTime = (Long) mContext.getAttribute(APP_PROXY_CALLS_TOTAL_TIME);
		vTime = new Long(vTime.longValue() + pTime);
		mContext.setAttribute(APP_PROXY_CALLS_TOTAL_TIME, vTime);
	}

	/**
	 * 
	 * @see javax.servlet.Servlet#getServletInfo()
	 */
	public String getServletInfo() {
		return "Proxy Servlet";
	}

	/**
	 * 
	 * @see javax.servlet.Servlet#destroy()
	 */
	public void destroy() {

	}
}